import sys
import os
import itertools
import threading
import queue
import time
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
                             QLabel, QLineEdit, QPushButton, QFileDialog, QComboBox,
                             QProgressBar, QTextEdit, QGroupBox, QSpinBox, QMessageBox)
from PyQt5.QtCore import Qt, pyqtSignal, QObject, QThread
import py7zr
import zipfile
import rarfile
import pyzipper


class WorkerSignals(QObject):
    update_progress = pyqtSignal(int, int, int)  # 当前进度, 总数, 百分比
    found_password = pyqtSignal(str)
    log_message = pyqtSignal(str)
    estimate_time = pyqtSignal(float)
    finished = pyqtSignal()


class ArchiveCracker(QThread):
    FORMATS = {
        '7z': ('*.7z', py7zr.is_7zfile),
        'ZIP': ('*.zip', zipfile.is_zipfile),
        'RAR': ('*.rar', rarfile.is_rarfile),
        'AES-ZIP': ('*.zip', zipfile.is_zipfile)
    }

    def __init__(self, archive_path, min_length, max_length, charset, output_dir=None,
                 thread_count=4, attack_type='bruteforce', dictionary_path=None, file_type='7z'):
        super().__init__()
        self.archive_path = archive_path
        self.min_length = min_length
        self.max_length = max_length
        self.charset = charset
        self.output_dir = output_dir
        self.thread_count = thread_count
        self.attack_type = attack_type
        self.dictionary_path = dictionary_path
        self.file_type = file_type
        self.signals = WorkerSignals()
        self._stop_event = threading.Event()
        self.password_queue = queue.Queue()
        self.total_passwords = self.calculate_total_passwords()
        self.extract_timeout_sec = 5
        self.progress = 0
        self.lock = threading.Lock()

    def calculate_total_passwords(self):
        """计算所有可能的密码数量"""
        if self.attack_type == 'dictionary':
            try:
                with open(self.dictionary_path, 'r', encoding='utf-8', errors='ignore') as f:
                    return sum(1 for line in f if line.strip())
            except:
                return 0
        else:
            return sum(len(self.charset) ** l for l in range(self.min_length, self.max_length + 1))

    def stop(self):
        self._stop_event.set()
        self.terminate()

    def password_generator(self):
        """密码生成器"""
        if self.attack_type == 'dictionary':
            with open(self.dictionary_path, 'r', encoding='utf-8', errors='ignore') as f:
                for line in f:
                    if self._stop_event.is_set():
                        return
                    password = line.strip()
                    if password:
                        yield password
        else:
            for length in range(self.min_length, self.max_length + 1):
                for pwd in itertools.product(self.charset, repeat=length):
                    if self._stop_event.is_set():
                        return
                    yield ''.join(pwd)

    def test_archive_password(self, password):
        """测试压缩包密码"""
        try:
            if self.file_type == '7z':
                with py7zr.SevenZipFile(self.archive_path, mode='r', password=password) as z:
                    z.list()
                    return True
            elif self.file_type == 'ZIP':
                with zipfile.ZipFile(self.archive_path, 'r') as z:
                    z.testzip()
                    return True
            elif self.file_type == 'AES-ZIP':
                with pyzipper.AESZipFile(self.archive_path) as z:
                    z.testzip()
                    return True
            elif self.file_type == 'RAR':
                with rarfile.RarFile(self.archive_path) as z:
                    if z.needs_password():
                        z.testrar()
                    return True
            return False
        except (py7zr.exceptions.PasswordRequired, zipfile.BadZipFile,
                rarfile.PasswordRequired, rarfile.BadRarFile, pyzipper.BadZipFile):
            return False
        except Exception as e:
            self.signals.log_message.emit(f"尝试密码 {password} 时出错: {str(e)}")
            return False

    def extract_archive(self, password):
        """解压压缩包"""
        try:
            if self.file_type == '7z':
                with py7zr.SevenZipFile(self.archive_path, mode='r', password=password) as z:
                    z.extractall(path=self.output_dir)
            elif self.file_type in ['ZIP', 'AES-ZIP']:
                if self.file_type == 'AES-ZIP':
                    with pyzipper.AESZipFile(self.archive_path) as z:
                        z.extractall(path=self.output_dir, pwd=password.encode('utf-8'))
                else:
                    with zipfile.ZipFile(self.archive_path, 'r') as z:
                        z.extractall(path=self.output_dir, pwd=password.encode('utf-8'))
            elif self.file_type == 'RAR':
                with rarfile.RarFile(self.archive_path) as z:
                    z.extractall(path=self.output_dir, pwd=password)
            return True
        except Exception as e:
            # self.signals.log_message.emit(f"解压失败: {str(e)}")
            return False

    def safe_extract(self, password):
        """带超时机制的安全解压方法"""

        def extract():
            try:
                extract_success[0] = self.extract_archive(password)
            except Exception as e:
                extract_error[0] = str(e)

        extract_success = [False]
        extract_error = [None]
        extract_thread = threading.Thread(target=extract)
        extract_thread.start()
        extract_thread.join(timeout=self.extract_timeout_sec)

        if extract_thread.is_alive():
            self.signals.log_message.emit(f"解压超时，可能密码 '{password}' 不正确")
            return False
        elif not extract_success[0]:
            self.signals.log_message.emit(f"密码: {password} ->解压失败: {extract_error[0] if extract_error[0] else '未知错误'}")
            return False
        return True

    def worker_task(self):
        """工作线程任务"""
        while not self._stop_event.is_set():
            try:
                password = self.password_queue.get_nowait()

                if self.test_archive_password(password):
                    if self.output_dir:
                        if not self.safe_extract(password):
                            continue

                    with self.lock:
                        self._stop_event.set()
                        self.signals.found_password.emit(password)
                        return

                with self.lock:
                    self.progress += 1
                    if self.progress % 100 == 0:
                        percent = int((self.progress / self.total_passwords) * 100)
                        self.signals.update_progress.emit(self.progress, self.total_passwords, percent)

            except queue.Empty:
                time.sleep(0.1)
                continue

    def run(self):
        """启动多线程破解"""
        try:
            attack_name = "字典攻击" if self.attack_type == 'dictionary' else "爆破攻击"
            self.signals.log_message.emit(f"开始{attack_name}，使用 {self.thread_count} 个线程...")
            start_time = time.time()

            # 填充密码队列
            for password in self.password_generator():
                self.password_queue.put(password)

            # 创建工作线程
            threads = []
            for _ in range(self.thread_count):
                t = threading.Thread(target=self.worker_task)
                t.start()
                threads.append(t)

            # 监控进度
            last_update_time = time.time()
            while not self._stop_event.is_set() and any(t.is_alive() for t in threads):
                time.sleep(0.1)

                # 更新进度和剩余时间
                current_time = time.time()
                if current_time - last_update_time >= 1:  # 每秒更新一次
                    with self.lock:
                        percent = int((self.progress / self.total_passwords) * 100)
                        self.signals.update_progress.emit(self.progress, self.total_passwords, percent)

                        elapsed = current_time - start_time
                        if self.progress > 0:
                            remaining = (self.total_passwords - self.progress) * (elapsed / self.progress)
                            self.signals.estimate_time.emit(remaining)

                        last_update_time = current_time

            # 清理线程
            for t in threads:
                t.join(timeout=0.1)

            if not self._stop_event.is_set():
                self.signals.update_progress.emit(self.total_passwords, self.total_passwords, 100)
                self.signals.log_message.emit(f"{attack_name}完成，未找到密码")

            self.signals.finished.emit()

        except Exception as e:
            self.signals.log_message.emit(f"发生错误: {str(e)}")
            self.signals.finished.emit()


class ArchiveCrackerGUI(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("多格式压缩包密码破解工具")
        self.setGeometry(100, 100, 800, 650)

        self.worker = None
        self.log_lines = 0
        self.max_log_lines = 200
        self.init_ui()

    def init_ui(self):
        main_widget = QWidget()
        layout = QVBoxLayout()

        # 文件选择部分
        file_group = QGroupBox("文件设置")
        file_layout = QVBoxLayout()

        # 压缩包文件选择
        hbox = QHBoxLayout()
        self.archive_label = QLabel("压缩包:")
        self.archive_path = QLineEdit()
        self.archive_browse = QPushButton("浏览...")
        self.archive_browse.clicked.connect(self.browse_archive)
        hbox.addWidget(self.archive_label)
        hbox.addWidget(self.archive_path)
        hbox.addWidget(self.archive_browse)
        file_layout.addLayout(hbox)

        # 压缩格式选择
        hbox = QHBoxLayout()
        hbox.addWidget(QLabel("压缩格式:"))
        self.file_type = QComboBox()
        self.file_type.addItems(['7z', 'ZIP', 'AES-ZIP', 'RAR'])
        hbox.addWidget(self.file_type)
        file_layout.addLayout(hbox)

        # 输出目录选择
        hbox = QHBoxLayout()
        self.output_label = QLabel("输出目录:")
        self.output_path = QLineEdit()
        self.output_browse = QPushButton("浏览...")
        self.output_browse.clicked.connect(self.browse_output)
        hbox.addWidget(self.output_label)
        hbox.addWidget(self.output_path)
        hbox.addWidget(self.output_browse)
        file_layout.addLayout(hbox)

        file_group.setLayout(file_layout)
        layout.addWidget(file_group)

        # 攻击方式选择
        attack_group = QGroupBox("攻击方式")
        attack_layout = QVBoxLayout()

        self.attack_type = QComboBox()
        self.attack_type.addItems(["纯爆破攻击", "字典攻击"])
        self.attack_type.currentIndexChanged.connect(self.update_attack_ui)
        attack_layout.addWidget(self.attack_type)

        # 线程数设置
        hbox = QHBoxLayout()
        hbox.addWidget(QLabel("线程数:"))
        self.thread_count = QSpinBox()
        self.thread_count.setRange(1, 16)
        self.thread_count.setValue(4)
        hbox.addWidget(self.thread_count)
        attack_layout.addLayout(hbox)

        # 纯爆破设置
        self.brute_frame = QWidget()
        brute_layout = QVBoxLayout()

        hbox = QHBoxLayout()
        hbox.addWidget(QLabel("最小长度:"))
        self.min_length = QSpinBox()
        self.min_length.setRange(1, 10)
        self.min_length.setValue(1)
        hbox.addWidget(self.min_length)

        hbox.addWidget(QLabel("最大长度:"))
        self.max_length = QSpinBox()
        self.max_length.setRange(1, 100)
        self.max_length.setValue(4)
        hbox.addWidget(self.max_length)
        brute_layout.addLayout(hbox)

        hbox = QHBoxLayout()
        hbox.addWidget(QLabel("字符集:"))
        self.charset = QLineEdit()
        self.charset.setText("0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")
        hbox.addWidget(self.charset)
        brute_layout.addLayout(hbox)

        self.brute_frame.setLayout(brute_layout)
        attack_layout.addWidget(self.brute_frame)

        # 字典攻击设置
        self.dict_frame = QWidget()
        dict_layout = QVBoxLayout()

        hbox = QHBoxLayout()
        self.dict_label = QLabel("字典文件:")
        self.dict_path = QLineEdit()
        self.dict_browse = QPushButton("浏览...")
        self.dict_browse.clicked.connect(self.browse_dict)
        hbox.addWidget(self.dict_label)
        hbox.addWidget(self.dict_path)
        hbox.addWidget(self.dict_browse)
        dict_layout.addLayout(hbox)

        self.dict_frame.setLayout(dict_layout)
        attack_layout.addWidget(self.dict_frame)
        self.dict_frame.hide()

        attack_group.setLayout(attack_layout)
        layout.addWidget(attack_group)

        # 进度信息
        progress_group = QGroupBox("进度信息")
        progress_layout = QVBoxLayout()

        # 进度条
        self.progress = QProgressBar()
        progress_layout.addWidget(self.progress)

        # 进度详情
        hbox = QHBoxLayout()
        self.progress_label = QLabel("进度: 0/0 (0%)")
        self.time_label = QLabel("预计剩余时间: --")
        hbox.addWidget(self.progress_label)
        hbox.addWidget(self.time_label)
        progress_layout.addLayout(hbox)

        # 组合数/密码数信息
        self.combination_label = QLabel("总组合数: --")
        progress_layout.addWidget(self.combination_label)

        progress_group.setLayout(progress_layout)
        layout.addWidget(progress_group)

        # 控制按钮
        hbox = QHBoxLayout()
        self.start_btn = QPushButton("开始")
        self.start_btn.clicked.connect(self.start_attack)
        self.stop_btn = QPushButton("停止")
        self.stop_btn.clicked.connect(self.stop_attack)
        self.stop_btn.setEnabled(False)
        hbox.addWidget(self.start_btn)
        hbox.addWidget(self.stop_btn)
        layout.addLayout(hbox)

        # 日志输出
        self.log = QTextEdit()
        self.log.setReadOnly(True)
        layout.addWidget(self.log)

        main_widget.setLayout(layout)
        self.setCentralWidget(main_widget)

    def update_attack_ui(self, index):
        if index == 0:  # 纯爆破
            self.brute_frame.show()
            self.dict_frame.hide()
            self.combination_label.setText("总组合数: --")
        else:  # 字典攻击
            self.brute_frame.hide()
            self.dict_frame.show()
            self.update_dict_count()

    def browse_archive(self):
        file_types = "压缩包文件 ({});;所有文件 (*.*)".format(
            " ".join([f"*{ext}" for ext in ['.7z', '.zip', '.rar']]))
        path, _ = QFileDialog.getOpenFileName(self, "选择压缩包", "", file_types)
        if path:
            self.archive_path.setText(path)

            # 自动检测文件类型
            if path.endswith('.7z'):
                self.file_type.setCurrentText('7z')
            elif path.endswith('.zip'):
                self.file_type.setCurrentText('ZIP')
            elif path.endswith('.rar'):
                self.file_type.setCurrentText('RAR')

    def browse_output(self):
        path = QFileDialog.getExistingDirectory(self, "选择输出目录")
        if path:
            self.output_path.setText(path)

    def browse_dict(self):
        path, _ = QFileDialog.getOpenFileName(self, "选择字典文件", "", "Text Files (*.txt);;All Files (*)")
        if path:
            self.dict_path.setText(path)
            self.update_dict_count()

    def update_dict_count(self):
        """更新字典密码数量显示"""
        if self.dict_path.text():
            try:
                with open(self.dict_path.text(), 'r', encoding='utf-8', errors='ignore') as f:
                    count = sum(1 for line in f if line.strip())
                self.combination_label.setText(f"字典密码数: {count}")
            except:
                self.combination_label.setText("字典密码数: 无法读取")
        else:
            self.combination_label.setText("字典密码数: --")

    def log_message(self, msg):
        """记录日志信息，限制最大行数"""
        self.log.append(msg)
        self.log_lines += 1

        # 保持日志行数不超过最大值
        if self.log_lines > self.max_log_lines:
            cursor = self.log.textCursor()
            cursor.movePosition(cursor.Start)
            cursor.select(cursor.LineUnderCursor)
            cursor.removeSelectedText()
            cursor.deleteChar()  # 删除换行符
            self.log_lines -= 1

    def update_progress(self, current, total, percent):
        """更新进度显示"""
        self.progress.setValue(percent)
        self.progress_label.setText(f"进度: {current}/{total} ({percent}%)")

        # 更新总组合数显示（如果是爆破模式）
        if self.attack_type.currentIndex() == 0 and self.combination_label.text() == "总组合数: --":
            self.combination_label.setText(f"总组合数: {total}")

    def update_estimate_time(self, seconds):
        """更新预估时间显示"""
        if seconds < 60:
            self.time_label.setText(f"预计剩余时间: {int(seconds)}秒")
        else:
            mins = int(seconds // 60)
            secs = int(seconds % 60)
            self.time_label.setText(f"预计剩余时间: {mins}分{secs}秒")

    def password_found(self, password):
        """密码找到时的处理"""
        self.log_message(f"成功找到密码: {password}")
        self.start_btn.setEnabled(True)
        self.stop_btn.setEnabled(False)

        # 弹出消息框
        msg = QMessageBox()
        msg.setIcon(QMessageBox.Information)
        msg.setWindowTitle("密码破解成功")
        msg.setText(f"成功找到密码: {password}")
        msg.setStandardButtons(QMessageBox.Ok)
        msg.exec_()

    def worker_finished(self):
        """工作线程完成时的处理"""
        if self.worker:
            self.worker.quit()
            self.worker.wait()
            self.worker = None
        self.start_btn.setEnabled(True)
        self.stop_btn.setEnabled(False)

    def start_attack(self):
        """开始攻击"""
        archive = self.archive_path.text()
        if not os.path.exists(archive):
            self.log_message("错误: 请选择有效的压缩包文件")
            return

        output_dir = self.output_path.text() if self.output_path.text() else None
        thread_count = self.thread_count.value()
        file_type = self.file_type.currentText()

        attack_type = 'dictionary' if self.attack_type.currentIndex() == 1 else 'bruteforce'

        # 重置UI状态
        self.progress.setValue(0)
        self.progress_label.setText("进度: 0/0 (0%)")
        self.time_label.setText("预计剩余时间: --")
        self.log.clear()
        self.log_lines = 0

        if attack_type == 'bruteforce':
            min_len = self.min_length.value()
            max_len = self.max_length.value()
            charset = self.charset.text()

            if min_len > max_len:
                self.log_message("错误: 最小长度不能大于最大长度")
                return

            if not charset:
                self.log_message("错误: 请指定字符集")
                return

            self.worker = ArchiveCracker(
                archive, min_len, max_len, charset, output_dir,
                thread_count, attack_type, file_type=file_type
            )

        else:  # 字典攻击
            dict_path = self.dict_path.text()
            if not os.path.exists(dict_path):
                self.log_message("错误: 请选择有效的字典文件")
                return

            self.worker = ArchiveCracker(
                archive, 0, 0, '', output_dir,
                thread_count, attack_type, dict_path, file_type
            )

        # 连接信号
        self.worker.signals.log_message.connect(self.log_message)
        self.worker.signals.update_progress.connect(self.update_progress)
        self.worker.signals.estimate_time.connect(self.update_estimate_time)
        self.worker.signals.found_password.connect(self.password_found)
        self.worker.signals.finished.connect(self.worker_finished)

        # 开始工作
        self.worker.start()
        self.start_btn.setEnabled(False)
        self.stop_btn.setEnabled(True)

        attack_name = "字典攻击" if attack_type == 'dictionary' else "爆破攻击"
        self.log_message(f"{attack_name}已开始，使用 {thread_count} 个线程...")
        self.log_message(f"目标文件: {archive}")
        self.log_message(f"文件类型: {file_type}")

    def stop_attack(self):
        """停止攻击"""
        if self.worker:
            self.worker.stop()
            self.log_message("正在停止攻击...")
            self.stop_btn.setEnabled(False)


if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = ArchiveCrackerGUI()
    window.show()
    sys.exit(app.exec_())